{ pkgs, lib, ... }:
let
  inherit (import ./ssh-keys.nix pkgs) snakeOilPrivateKey snakeOilPublicKey;

  mkNode = vlan: id: {
    virtualisation.vlans = [ vlan ];
    networking = {
      useDHCP = false;
      useNetworkd = true;
    };

    systemd.network = {
      enable = true;

      networks."10-eth${toString vlan}" = {
        matchConfig.Name = "eth${toString vlan}";
        linkConfig.RequiredForOnline = "no";
        networkConfig = {
          Address = "192.168.${toString vlan}.${toString id}/24";
          IPv4Forwarding = "yes";
          IPv6Forwarding = "yes";
        };
      };
    };
  };
in
{
  name = "systemd-networkd-vrf";
  meta.maintainers = with lib.maintainers; [ ma27 ];

  nodes = {
    client =
      { pkgs, ... }:
      {
        virtualisation.vlans = [
          1
          2
        ];

        networking = {
          useDHCP = false;
          useNetworkd = true;
          firewall.checkReversePath = "loose";
        };

        systemd.network = {
          enable = true;

          netdevs."10-vrf1" = {
            netdevConfig = {
              Kind = "vrf";
              Name = "vrf1";
              MTUBytes = "1300";
            };
            vrfConfig.Table = 23;
          };
          netdevs."10-vrf2" = {
            netdevConfig = {
              Kind = "vrf";
              Name = "vrf2";
              MTUBytes = "1300";
            };
            vrfConfig.Table = 42;
          };

          networks."10-vrf1" = {
            matchConfig.Name = "vrf1";
            networkConfig.IPv4Forwarding = "yes";
            networkConfig.IPv6Forwarding = "yes";
            routes = [
              {
                Destination = "192.168.1.2";
                Metric = 100;
              }
            ];
          };
          networks."10-vrf2" = {
            matchConfig.Name = "vrf2";
            networkConfig.IPv4Forwarding = "yes";
            networkConfig.IPv6Forwarding = "yes";
            routes = [
              {
                Destination = "192.168.2.3";
                Metric = 100;
              }
            ];
          };

          networks."10-eth1" = {
            matchConfig.Name = "eth1";
            linkConfig.RequiredForOnline = "no";
            networkConfig = {
              VRF = "vrf1";
              Address = "192.168.1.1/24";
              IPv4Forwarding = "yes";
              IPv6Forwarding = "yes";
            };
          };
          networks."10-eth2" = {
            matchConfig.Name = "eth2";
            linkConfig.RequiredForOnline = "no";
            networkConfig = {
              VRF = "vrf2";
              Address = "192.168.2.1/24";
              IPv4Forwarding = "yes";
              IPv6Forwarding = "yes";
            };
          };
        };
      };

    node1 = lib.mkMerge [
      (mkNode 1 2)
      {
        services.openssh.enable = true;
        users.users.root.openssh.authorizedKeys.keys = [ snakeOilPublicKey ];
      }
    ];

    node2 = mkNode 2 3;
    node3 = mkNode 2 4;
  };

  testScript = ''
    import json

    def compare(raw_json, to_compare):
        data = json.loads(raw_json)
        assert len(raw_json) >= len(to_compare)
        for i, row in enumerate(to_compare):
            actual = data[i]
            assert len(row.keys()) > 0
            for key, value in row.items():
                assert value == actual[key], f"""
                  In entry {i}, value {key}: got: {actual[key]}, expected {value}
                """


    start_all()

    client.wait_for_unit("network.target")
    node1.wait_for_unit("network.target")
    node2.wait_for_unit("network.target")
    node3.wait_for_unit("network.target")

    # Check that networkd properly configures the main routing table
    # and the routing tables for the VRF.
    with subtest("check vrf routing tables"):
        compare(
            client.succeed("ip --json -4 route list"),
            [
                {"dst": "192.168.1.2", "dev": "vrf1", "metric": 100},
                {"dst": "192.168.2.3", "dev": "vrf2", "metric": 100}
            ]
        )
        compare(
            client.succeed("ip --json -4 route list table 23"),
            [
                {"dst": "192.168.1.0/24", "dev": "eth1", "prefsrc": "192.168.1.1"},
                {"type": "local", "dst": "192.168.1.1", "dev": "eth1", "prefsrc": "192.168.1.1"},
                {"type": "broadcast", "dev": "eth1", "prefsrc": "192.168.1.1", "dst": "192.168.1.255"}
            ]
        )
        compare(
            client.succeed("ip --json -4 route list table 42"),
            [
                {"dst": "192.168.2.0/24", "dev": "eth2", "prefsrc": "192.168.2.1"},
                {"type": "local", "dst": "192.168.2.1", "dev": "eth2", "prefsrc": "192.168.2.1"},
                {"type": "broadcast", "dev": "eth2", "prefsrc": "192.168.2.1", "dst": "192.168.2.255"}
            ]
        )

    # Ensure that other nodes are reachable via ICMP through the VRF.
    with subtest("icmp through vrf works"):
        client.succeed("ping -c5 192.168.1.2")
        client.succeed("ping -c5 192.168.2.3")

    # Test whether TCP through a VRF IP is possible.
    with subtest("tcp traffic through vrf works"):
        node1.wait_for_open_port(22)
        client.succeed(
            "cat ${snakeOilPrivateKey} > privkey.snakeoil"
        )
        client.succeed("chmod 600 privkey.snakeoil")
        client.succeed(
            "ulimit -l 2048; ip vrf exec vrf1 ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i privkey.snakeoil root@192.168.1.2 true"
        )

    # Only configured routes through the VRF from the main routing table should
    # work. Additional IPs are only reachable when binding to the vrf interface.
    with subtest("only routes from main routing table work by default"):
        client.fail("ping -c5 192.168.2.4")
        client.succeed("ping -I vrf2 -c5 192.168.2.4")

    client.shutdown()
    node1.shutdown()
    node2.shutdown()
    node3.shutdown()
  '';
}
